{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "# Transformer的交通流量预测 完整数据代码 整合了时间天气等多方面信息\n",
    "\n",
    "import datetime\n",
    "from torch.utils.data import DataLoader\n",
    "from dataset import TrafficDataset\n",
    "from model import TransAm, LSTMModel, lstm_init_weights, Transformer\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "import math\n",
    "from utils import PlotUtils\n",
    "from constant import Constant\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import os\n",
    "\n",
    "\n",
    "def test(args, model, data_loader, criterion):\n",
    "    model.eval()\n",
    "    total_loss = 0.\n",
    "    predictions = torch.Tensor(0)\n",
    "    truth = torch.Tensor(0)\n",
    "    with torch.no_grad():\n",
    "        for data, label, tgt_in in data_loader:\n",
    "            data, label = data.to(args.device).float(), label.to(args.device).float()\n",
    "            # label = label.reshape((label.shape[0], label.shape[1], -1))\n",
    "            # label = label.permute((1, 0, 2))\n",
    "            if args.have_decoder:\n",
    "                tgt_in = tgt_in.to(args.device).float()\n",
    "                output = model(data, tgt_in)\n",
    "            else:\n",
    "                output = model(data)\n",
    "            total_loss += criterion(output, label).item()\n",
    "            predictions = torch.cat((predictions, output[:, -Constant.OUTPUT_WINDOW].view(-1).cpu()), 0)\n",
    "            truth = torch.cat((truth, label[:, -Constant.OUTPUT_WINDOW].view(-1).cpu()), 0)\n",
    "\n",
    "        loss = total_loss / len(data_loader)\n",
    "        tqdm.write(f'valid loss : {loss:.5f}, ppl:{math.exp(loss):.5f}')\n",
    "        tqdm.write(\"-\" * 100)\n",
    "    return truth, predictions, loss, math.exp(loss)\n",
    "\n",
    "\n",
    "def train(args, model, data_loader, criterion, optimizer, epoch, writer):\n",
    "    model.train()\n",
    "    batch_len = len(data_loader)\n",
    "\n",
    "    process_bar = tqdm(total=batch_len)\n",
    "    for idx, (data, label, tgt_in) in enumerate(data_loader):\n",
    "        process_bar.update(1)\n",
    "        data, label = data.to(args.device).float(), label.to(args.device).float()\n",
    "        # label = label.reshape((label.shape[0], label.shape[1], -1))\n",
    "        # label = label.permute((1, 0, 2))\n",
    "        optimizer.zero_grad()\n",
    "        if args.have_decoder:\n",
    "            tgt_in = tgt_in.to(args.device).float()\n",
    "            output = model(data, tgt_in)\n",
    "        else:\n",
    "            output = model(data)\n",
    "        loss = criterion(output, label)\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.7)\n",
    "        optimizer.step()\n",
    "\n",
    "        writer.add_scalar(\"train_loss\", loss.item(), epoch * batch_len + idx)\n",
    "        writer.add_scalar(\"train_ppl\", math.exp(loss.item()), epoch * batch_len + idx)\n",
    "\n",
    "        if idx % 10 == 0:\n",
    "            tqdm.write(f'[epoch:{epoch:>3d}] loss:{loss.item():.7f} , ppl:{math.exp(loss.item())} .')\n",
    "\n",
    "\n",
    "def get_model(args):\n",
    "    \"\"\"模型实例化\"\"\"\n",
    "    feature_choose = Constant.FEATURES_CHOOSE\n",
    "    n_inputs = len(Constant.FEATURES_MAP[feature_choose]) + 1\n",
    "\n",
    "    model_name = Constant.MODELS[Constant.MODEL_CHOOSE]\n",
    "    print(f'you are using {model_name}...')\n",
    "    args.model = model_name\n",
    "    if model_name == 'transformer_encoder':\n",
    "        return TransAm(input_feature_size=n_inputs).to(args.device)\n",
    "    elif model_name == 'LSTM':\n",
    "        return LSTMModel(n_input=n_inputs).to(args.device).apply(lstm_init_weights)\n",
    "    elif model_name == 'transformer':\n",
    "        args.have_decoder = True\n",
    "        return Transformer(n_encoder_inputs=n_inputs, n_decoder_inputs=n_inputs).to(args.device)\n",
    "    else:\n",
    "        raise ValueError\n",
    "\n",
    "\n",
    "def main(main_args):\n",
    "    # 获取模型\n",
    "    model = get_model(main_args)\n",
    "    print(main_args.have_decoder)\n",
    "\n",
    "    print(f'you are using features: {Constant.FEATURES_MAP[Constant.FEATURES_CHOOSE]}'\n",
    "          f'to predict {Constant.TARGET}...')\n",
    "    train_dataset = TrafficDataset(train=True)\n",
    "    train_dataset_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=args.shuffle)\n",
    "    test_dataset = TrafficDataset(train=False)\n",
    "    test_dataset_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
    "\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learn_rate)\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)\n",
    "\n",
    "    writer = SummaryWriter(args.tensorboard)\n",
    "\n",
    "    best = None\n",
    "\n",
    "    for epoch in range(args.epochs):\n",
    "        # train\n",
    "        train(main_args, model, train_dataset_loader, criterion, optimizer, epoch, writer)\n",
    "        # validation\n",
    "        truth, predictions, valid_loss, valid_ppl = test(args, model, test_dataset_loader, criterion)\n",
    "\n",
    "        if best is None:\n",
    "            best = [valid_loss, valid_ppl]\n",
    "        if valid_loss <= best[0]:\n",
    "            best[0] = valid_loss\n",
    "            best[1] = valid_ppl\n",
    "            tqdm.write(f'model saving...')\n",
    "            model_save_path = os.path.join(args.backup, f'{args.model}_best_{valid_loss:.5f}.pth')\n",
    "            torch.save(model.state_dict(), model_save_path)\n",
    "\n",
    "        writer.add_scalar(\"valid_loss\", valid_loss, epoch)\n",
    "        writer.add_scalar(\"valid_ppl\", valid_ppl, epoch)\n",
    "        PlotUtils.plot(truth, predictions, epoch)\n",
    "        # scheduler\n",
    "        scheduler.step()\n",
    "\n",
    "#  基于 transformer、LSTM的交通流量预测 完整数据代码 超全项目 本科硕士毕设卓卓有余\n",
    "# loss 下降很稳定 结果很好\n",
    "# 一键搭配环境可以参考：https://pythonwfyyqq.blog.csdn.net/article/details/121321372\n",
    "if __name__ == \"__main__\":\n",
    "    now_time = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--epochs\", type=int, default=100)\n",
    "    parser.add_argument(\"--learn_rate\", \"-lr\", type=float, default=0.005)\n",
    "    parser.add_argument(\"--device\", type=str)\n",
    "    parser.add_argument(\"--shuffle\", type=bool, default=True)\n",
    "    parser.add_argument(\"--batch_size\", type=int, default=10)\n",
    "    parser.add_argument(\"--have_decoder\", type=bool, default=False)\n",
    "    parser.add_argument(\"--model\", type=str)\n",
    "    parser.add_argument(\"--backup\", type=str, default='backup/',\n",
    "                        help=\"the directory to save model parameter\")\n",
    "    parser.add_argument(\"--tensorboard\", type=str, default=os.path.join(\"tensorboard\", now_time),\n",
    "                        help=\"the directory to save tensorboard data\")\n",
    "\n",
    "    args = parser.parse_args()\n",
    "    args.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    print(f\"You are using {args.device}......\")\n",
    "\n",
    "    with open(\"tensorboard/tensorboard.txt\", 'w') as handle:\n",
    "        handle.write(f'tensorboard --logdir={args.tensorboard}')\n",
    "\n",
    "    main(args)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
